#!/usr/env/bin python
# -*- coding: utf-8 -*-

# @Time    : 2021-07-05 9:53 上午
# @Author  : yangdy
# @File    : 2021-07-05.726.原子的数量
# @Software: PyCharmCE

from tools.my_methods import *


class Solution:
    def countOfAtoms(self, formula: str) -> str:
        stack = [defaultdict(int), ]
        # uppers = ''.join(chr(ord('A')+i) for i in range(26))
        # lowers = ''.join(chr(ord('a')+i) for i in range(26))
        i = 0
        x, d = '', 0
        while i < len(formula):
            si = formula[i]
            if si in "0123456789":
                d = d * 10 + ord(si) - ord('0')
                i += 1
            elif si.isupper():
                if x != '':
                    if d == 0:
                        d = 1
                    stack[-1][x] += d
                    x, d = '', 0
                x = si
                i += 1
                while i < len(formula) and formula[i].islower():
                    x += formula[i]
                    i += 1
            elif si == '(':
                if x != '':
                    if d == 0:
                        d = 1
                    stack[-1][x] += d
                    x, d = '', 0
                stack.append(defaultdict(int))
                i += 1
            elif si == ')':
                if x != '':
                    if d == 0:
                        d = 1
                    stack[-1][x] += d
                    x, d = '', 0
                mm = stack.pop()
                dd = 0
                i += 1
                while i < len(formula) and formula[i].isdigit():
                    dd = dd*10+ord(formula[i])-ord('0')
                    i += 1
                if dd == 0:
                    dd = 1
                for k in mm.keys():
                    stack[-1][k] += mm[k]*dd
        if x != '':
            if d == 0:
                d = 1
            stack[-1][x] += d
        mm = stack.pop()
        keys = [k for k in mm.keys()]
        keys.sort()
        return ''.join(f'{k}{mm[k] if mm[k]>1 else ""}' for k in keys)


questions = [
    "H2O",
    "Mg(OH)2",
    "K4(ON(SO3)2)2"
]

work(questions, lambda qi: Solution().countOfAtoms(qi))
